import csv

from utils import infer_pvalue, csv_to_fasta, retrive_esm1b_embedding, mutate_single_seq_ECs, compute_esm_distance

train_data = '/state/partition/wzzheng/clean/data/train_valid_split/split100/split100_train_split_0'
test_data = '/state/partition/wzzheng/clean/data/datasets/new'
esm_dir= '/state/partition/wzzheng/clean/data/train_valid_split/split100/esm_data'
model_name = 'rnn'
model_path = '/state/partition/wzzheng/clean/model/rnn_split100_train_split_0.pth'
results_dir = '/state/partition/wzzheng/clean/results'


def check_overlap(train_file, test_file):
    with open(train_file, 'r') as f:
        train_data = set(row[0] for row in csv.reader(f))

    with open(test_file, 'r') as f:
        test_data = set(row[0] for row in csv.reader(f))

    overlap = train_data & test_data  # Find the intersection of the two sets

    if overlap:
        print(f"Found {len(overlap)} overlapping IDs:")
        for id in overlap:
            print(id)
    else:
        print("No overlapping IDs found.")

# check_overlap(train_data, test_data)


csv_to_fasta(test_data + ".csv", test_data + ".fasta")
retrive_esm1b_embedding(fasta_name=test_data)

infer_pvalue(train_data=train_data, test_data=test_data, esm_dir=esm_dir,model_name=model_name, model_path=model_path, results_dir=results_dir)